import logging
import threading
import time
from iotdb.Session import Session


class SessionPool:
    def __init__(self, host, port, user, password, max_size=10, min_size=3, connection_timeout=10000):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.max_size = max_size
        self.min_size = min_size
        self.connection_timeout = connection_timeout  # 毫秒
        self.pool = []
        self.lock = threading.Lock()
        self.is_closed = False
        self.health_check_interval = 60  # 秒
        self._initialize_pool()
        self._start_health_check()

    def _initialize_pool(self):
        """初始化连接池，创建最小数量的连接"""
        with self.lock:
            for _ in range(self.min_size):
                try:
                    session = self._create_session()
                    self.pool.append(session)
                except Exception as e:
                    logging.error(f"创建IoTDB会话失败: {e}")

    def _create_session(self):
        """创建一个新的IoTDB会话并打开连接"""
        session = Session(self.host, self.port, self.user, self.password, connection_timeout_ms=self.connection_timeout)
        session.open(False)
        return session

    def _check_session(self, session):
        """检查会话是否有效"""
        try:
            # 尝试一个简单的查询来验证连接有效性
            session.execute_query_statement("SHOW TIMESERIES")
            return True
        except Exception:
            return False

    def _health_check(self):
        """定期检查连接池中的所有连接"""
        while not self.is_closed:
            try:
                with self.lock:
                    # 过滤出无效连接
                    valid_sessions = []
                    for session in self.pool:
                        if self._check_session(session):
                            valid_sessions.append(session)
                        else:
                            try:
                                session.close()
                            except:
                                pass  # 忽略关闭错误
                    
                    # 更新连接池
                    self.pool = valid_sessions
                    
                    # 如果连接数低于最小值，创建新连接
                    while len(self.pool) < self.min_size:
                        try:
                            session = self._create_session()
                            self.pool.append(session)
                        except Exception as e:
                            logging.error(f"健康检查: 创建新会话失败: {e}")
                            break
                    
                    logging.debug(f"IoTDB连接池健康检查: 当前活跃连接 {len(self.pool)}/{self.max_size}")
            except Exception as e:
                logging.error(f"IoTDB连接池健康检查失败: {e}")
            
            # 等待下一次检查
            time.sleep(self.health_check_interval)

    def _start_health_check(self):
        """启动连接池健康检查线程"""
        health_check_thread = threading.Thread(target=self._health_check, daemon=True)
        health_check_thread.start()

    def get_session(self, max_retries=3):
        """获取一个会话，如果池为空则创建新会话"""
        for attempt in range(max_retries):
            with self.lock:
                if self.is_closed:
                    raise Exception("会话池已关闭")
                
                if self.pool:
                    session = self.pool.pop(0)
                    # 验证会话有效性
                    if self._check_session(session):
                        return session
                    else:
                        # 会话无效，关闭并重试
                        try:
                            session.close()
                        except:
                            pass
                
                # 池为空或会话无效，创建新会话
                try:
                    return self._create_session()
                except Exception as e:
                    if attempt == max_retries - 1:
                        raise Exception(f"无法创建IoTDB会话: {e}")
                    time.sleep(1)  # 重试前等待
        
        raise Exception("无法获取有效的IoTDB会话")

    def put_session(self, session):
        """归还会话到连接池，如果池已满则关闭会话"""
        if session is None:
            return
            
        with self.lock:
            if self.is_closed:
                try:
                    session.close()
                except:
                    pass
                return
                
            # 检查会话是否有效
            if not self._check_session(session):
                try:
                    session.close()
                except:
                    pass
                return
                
            # 如果池未满则归还会话
            if len(self.pool) < self.max_size:
                self.pool.append(session)
            else:
                try:
                    session.close()
                except:
                    pass

    def close_all(self):
        """关闭连接池中的所有会话"""
        with self.lock:
            self.is_closed = True
            for session in self.pool:
                try:
                    session.close()
                except Exception as e:
                    logging.error(f"关闭IoTDB会话失败: {e}")
            self.pool.clear()
            logging.info("IoTDB会话池已关闭")